{-# LANGUAGE TupleSections #-}

module Expr (T(..), evalP, ElemFun(..), elemFun, prob) where

import qualified Iv
import Interlude
import Text.Parsec hiding ((<|>),many)
import Text.Parsec.Token
import Text.Parsec.Language (haskellDef)
import Text.Parsec.Expr

-- * Signatures

data T c =
   Const c
 | Pi
 | Var Int
 | Add (T c) (T c)
 | Mul (T c) (T c)
 | Neg (T c)
 | Pow (T c) Integer
 | Elem ElemFun (T c) deriving Show

data ElemFun = Exp | Sqrt | Atan | Recip deriving (Show, Enum, Bounded)

-- * Evaluation

eval :: Fractional a => (ElemFun -> a -> a) -> (Int -> a) -> (c -> a) -> T c -> a
eval elm var cnst = evl where
  evl (Const s)  = cnst s
  evl (Add t u)  = evl t + evl u
  evl (Mul t u)  = evl t * evl u
  evl (Neg t)    = - evl t
  evl (Pow t n)  = evl t ^ n
  evl (Var i)    = var i
  evl (Elem f t) = elm f $ evl t

evalP :: Fractional a => (Int -> a) -> (c -> a) -> T c -> a
evalP = eval (error "Expr.evalP: elementary function")

elemFun :: Floating a => ElemFun -> a -> a
elemFun Exp = exp
elemFun Sqrt = sqrt
elemFun Atan = atan
elemFun Recip = recip

-- * Parsing

l :: TokenParser ()
l = makeTokenParser haskellDef { reservedNames = "Pi" : map show (enumFrom (minBound::ElemFun))}

term :: Fractional c => Bool -> [(String, Iv.T c)] -> Parser (T c)
term np x = parens l (expr np x) <|> brackets l (expr np x)
 <|> (Const <$> (try (realToFrac <$> float l) <|> (fromInteger <$> natural l <* optional (reservedOp l "."))))
 <|> (if np then Pi <$ reserved l "Pi" else mzero)
 <|> (identifier l >>= \v -> maybe (parserFail $ "Undeclared variable: " ++ v)
                   (return . Var) (elemIndex v $ map fst x))

expr :: Fractional c => Bool -> [(String, Iv.T c)] -> Parser (T c)
expr np x = buildExpressionParser
  [ [ Postfix $ flip Pow <$> ((reservedOp l "^" <|> reservedOp l "**") *> natural l) ]
  , [ Prefix $ Neg        <$ reservedOp l "-"
    , Prefix $ id         <$ reservedOp l "+"
    , Prefix $ Elem Recip <$ reservedOp l "/" ]
  , if np then map (\f -> Prefix $ Elem f <$ reserved l (show f)) (enumFrom minBound) else []
  , [ Infix (Mul                  <$ reservedOp l "*") AssocLeft
    , Infix ((. Elem Recip) . Mul <$ reservedOp l "/") AssocLeft ]
  , [ Infix ( Add                 <$ reservedOp l "+") AssocLeft
    , Infix ((. Neg) . Add        <$ reservedOp l "-") AssocLeft ]
  ] (term np x)

-- Note: Variable bounds cannot contain elementary functions
-- (e.g. Sqrt(2))

box :: Fractional a => Parser [(String, Iv.T a)]
box = many $ liftA3 (\a x b -> (x, Iv.mk a b))
  (try (evalP undefined id <$> expr False [] <* reservedOp l "<="))
  (identifier l <* reservedOp l "<=")
  (evalP undefined id <$> expr False [] <* reservedOp l "->")

disj :: Fractional a => Bool -> [(String, Iv.T a)] -> Parser [Expr.T a]
disj np x = sepBy1 (expr np x) (reservedOp l "\\/")

prob :: Fractional a => Bool -> Parser ([Iv.T a], [Expr.T a])
prob np = whiteSpace l >> box >>= \x -> (map snd x,) <$> disj np x <* eof
